import os
import sys
import time
import json
import concurrent.futures

import requests
import bs4
from bs4 import BeautifulSoup
import pdfkit

config = {}
oj_url = ''
cookies = {}
config_file = os.path.join(os.path.split(os.path.realpath(__file__))[0], 'config.json')
pdfkit_config = None
debug = True


class ThreadPoolExecutorPlus(concurrent.futures.ThreadPoolExecutor):
    def advance_map(self, func, *iterables, timeout=None, max_workers=None, max_try_times=None,
                    exception_func=None, success_func=None, error_func=None, timeout_func=None, callback_func=None):
        """
        :param func: 操作于可迭代对象上的函数
        :param iterables: 一个或多个可迭代对象
        :param timeout: 超时时间，不指定则不限时
        :param max_workers: 线程池中的线程数
        :param max_try_times: 最多重试次数，默认为 1
        :param exception_func: func(e, *args) 每次重试时出现异常时调用，不指定则为抛出错误
        :param success_func: func(*args) 成功运行后时调用
        :param error_func: func(e, *args) 超过了最大重试次数，仍为成功时调用
        :param timeout_func: func(*args) 超时时调用
        :param callback_func: func(future) 添加为每个 Future 对象的回调函数
        :return: 返回一个将 function 应用于 iterable 中每一项并输出其结果的迭代器
        :rtype: Iterable
        """
        if max_workers is not None:
            self.__init__(max_workers=max_workers)
        if timeout is not None:
            end_time = timeout + time.monotonic()
        if max_try_times is None:
            max_try_times = 1
        if max_try_times < 1:
            raise ValueError("max_try_times must > 1")

        def _run_warp(fn, /, *args):
            try_times = 0
            last_exception = None
            while try_times < max_try_times:
                try:
                    return fn(*args)
                except Exception as e:
                    last_exception = e
                    if exception_func is not None:
                        exception_func(e, *args)
                    else:
                        # raise e
                        pass
                finally:
                    try_times += 1
            raise last_exception

        fs = [(self.submit(_run_warp, func, *args), args) for args in zip(*iterables)]
        if callback_func is not None:
            for future in fs:
                future[0].add_done_callback(callback_func)

        def result_iterator():
            try:
                fs.reverse()
                while fs:
                    _future, _args = fs.pop()
                    try:
                        if timeout is None:
                            yield _future.result()
                        else:
                            yield _future.result(end_time - time.monotonic())
                    except concurrent.futures.TimeoutError as e:
                        if timeout_func is not None:
                            yield timeout_func(*_args)
                        elif error_func is not None:
                            yield error_func(e, *_args)
                        else:
                            raise e
                    except Exception as e:
                        if error_func is not None:
                            yield error_func(e, *_args)
                        else:
                            raise e
                    else:
                        if success_func is not None:
                            success_func(*_args)
            finally:
                for _future in fs:
                    _future[0].cancel()

        return result_iterator()


class WZUOJProblem(object):
    def __init__(self, problem_id, title, url, volume_info):
        self.id = problem_id
        self.title = title
        self.url = url
        self.volume_info = volume_info

    def get_format_path(self, judge_id, path_type):
        format_path = config[path_type].format(id=self.id, title=self.title,
                                               volume_name=self.volume_info['name'],
                                               volume_id=self.volume_info['id'],
                                               course_name=self.volume_info['course_name'],
                                               judge_id=judge_id)
        os.makedirs(os.path.dirname(format_path), exist_ok=True)
        return format_path

    def get_judge_list(self):
        r = requests.get(
            oj_url + '/judge/course/{0}/judgelist/?problem={1}&userprofile={2}'.format(self.volume_info['course_id'],
                                                                                       self.id, config['user_profile']),
            cookies=cookies)

        soup = BeautifulSoup(r.text, 'lxml')
        table = soup.find('table')
        table_list = filter(lambda x: type(x) == bs4.element.Tag, table.tbody.children)
        return_list = []
        for item in table_list:
            item_contents = list(filter(lambda x: type(x) == bs4.element.Tag, item.contents))
            return_list.append({
                'id': str(item_contents[0].a.string).strip(),
                'status': item_contents[5].get_text().strip()
                # others is unnecessary
            })
        return return_list

    def download_pdf(self):
        try:
            ac_judge_list = filter(lambda x: x['status'] == 'AC', self.get_judge_list())
            last_ac_judge = next(ac_judge_list)
            if not os.path.isfile(self.get_format_path(last_ac_judge['id'], 'pdf_path')):
                pdfkit.from_url(
                    oj_url + '/judge/{}/course/{}/print_exp/'.format(last_ac_judge['id'],
                                                                     self.volume_info['course_id']),
                    self.get_format_path(last_ac_judge['id'], 'pdf_path'),
                    options=get_pdfkit_options(), configuration=pdfkit_config)
            return None
        except Exception as e:
            return e

    def download_source_code(self):
        try:
            ac_judge_list = filter(lambda x: x['status'] == 'AC', self.get_judge_list())
            last_ac_judge = next(ac_judge_list)
            r = requests.get(oj_url + '/judge/{}/course/{}/'.format(last_ac_judge['id'], self.volume_info['course_id']),
                             cookies=cookies)
            soup = BeautifulSoup(r.text, 'lxml')
            pre_list = soup.find_all('pre')

            def is_source_code_pre(pre):
                try:
                    return pre.parent.parent.td.string == '源代码'
                except AttributeError:
                    return False

            source_code = str(next(filter(is_source_code_pre, pre_list)).string).replace('\r', '')
            source_code_path = self.get_format_path(last_ac_judge['id'], 'sc_path')
            with open(source_code_path, 'w', encoding='utf8') as fp:
                fp.write(source_code)
            return None
        except Exception as e:
            return e


def get_pdfkit_options():
    if debug:
        return {'cookie': list(cookies.items())}
    else:
        return {'quiet': '', 'cookie': list(cookies.items())}


def get_course_detail(course_id):
    r = requests.get(oj_url + '/course/' + str(course_id) + '/detail/', cookies=cookies)
    soup = BeautifulSoup(r.text, 'lxml')
    course_name = str(soup.find('a', class_='brand').string)
    table = soup.find('table')
    table_list = filter(lambda x: type(x) == bs4.element.Tag, table.tbody.children)
    return_list = []
    for item in table_list:
        item_contents = list(filter(lambda x: type(x) == bs4.element.Tag, item.contents))
        return_list.append({
            'course_name': course_name,
            'course_id': course_id,
            'id': str(item_contents[0].string).strip(),
            'name': str(item_contents[1].a.string).strip(),
            'url': str(item_contents[1].a['href']),
            'level': int(item_contents[2].string),
            'quantity': int(item_contents[3].string)
        })
    return return_list


def get_problem_list(volume_info, problem_status=None):
    r = requests.get(oj_url + volume_info['url'], cookies=cookies)
    soup = BeautifulSoup(r.text, 'lxml')
    table = soup.find('div', id='problemStatus').table.tbody
    table_list = filter(lambda x: type(x) == bs4.element.Tag, table.children)
    return_list = []
    for item in table_list:
        item_contents = list(filter(lambda x: type(x) == bs4.element.Tag, item.contents))
        if problem_status is None or item_contents[0].get_text().strip() == problem_status:
            return_list.append(WZUOJProblem(problem_id=int(item_contents[1].string),
                                            title=str(item_contents[2].a.string).strip(),
                                            url=item_contents[2].a['href'],
                                            volume_info=volume_info))
    return return_list


def print_course_detail(volume_list):
    print('\n' + volume_list[0]['course_name'] + '\n\n[0] ALL')
    for index, item in enumerate(volume_list, 1):
        print('[{0}] {1} {2}'.format(index, item['id'], item['name']))


def load_config():
    global config
    global oj_url
    global cookies
    global pdfkit_config
    try:
        with open(config_file, 'r', encoding='utf8') as fp:
            config = json.load(fp)
    except json.decoder.JSONDecodeError:
        print('配置文件读取错误！')
        sys.exit()
    if config['use_webvpn']:
        oj_url = config['wzu_webvpn_oj_url']
        cookies = config['webvpn_cookies']
    else:
        oj_url = config['wzu_oj_url']
        cookies = config['cookies']
    if config.get('wkhtmltopdf_path', None):
        pdfkit_config = pdfkit.configuration(wkhtmltopdf=config['wkhtmltopdf_path'])


def main():
    load_config()
    volume_list = get_course_detail(input('请输入课程代码：'))
    print_course_detail(volume_list)
    while True:
        print('\np [num] 下载PDF，s [num] 下载源代码.')
        choose = input('请输入命令：')
        choose_list = choose.split()
        if len(choose_list) == 2 and choose_list[1].isdecimal() and 0 <= int(choose_list[1]) <= len(volume_list):
            if not (choose_list[0] == 'p' or choose_list[0] == 's'):
                continue

            def downloader_func(problem):
                if choose_list[0] == 'p':
                    return problem.download_pdf()
                elif choose_list[0] == 's':
                    return problem.download_source_code()

            choose = int(choose_list[1])
            if choose == 0:
                problem_list = []
                for i in range(len(volume_list)):
                    problem_list += get_problem_list(volume_list[i], 'AC')
            else:
                problem_list = get_problem_list(volume_list[choose - 1], 'AC')

            if config['multithreading']:
                with ThreadPoolExecutorPlus() as executor:
                    executor.advance_map(downloader_func, problem_list,
                                         timeout=config.get('timeout', None),
                                         max_workers=config.get('max_workers', None),
                                         max_try_times=config.get('max_try_times', None))
            else:
                for index, item in enumerate(problem_list, 1):
                    time.sleep(config.get('sleep', 0))
                    downloader_return = downloader_func(item)
                    if downloader_return is None:
                        print("[{}/{}]Success!".format(index, len(problem_list)))
                    else:
                        print("[{}/{}]Failed! {}".format(index, len(problem_list), downloader_return))
            break
        else:
            print('输入错误！请重新输入.')


if __name__ == '__main__':
    main()
